import cv2
import numpy as np


def overlay_logo(background, logo_path, position=(0, 0), color=(255, 255, 255)):
    """
    将logo叠加到背景图像上

    Args:
        background: 背景图像
        logo_path: logo文件路径
        position: 叠加位置 (x, y)
        color: logo颜色 (B, G, R)
    """
    # 读取logo
    logo = cv2.imread(logo_path, cv2.IMREAD_UNCHANGED)
    if logo is None:
        print("Error: Could not load logo")
        return background

    # 分离通道
    logo_bgr = logo[:, :, :3]
    logo_alpha = logo[:, :, 3]

    # 创建彩色logo
    colored_logo = np.zeros_like(logo_bgr)
    colored_logo[:] = color  # 设置颜色

    # 获取背景的ROI区域
    x, y = position
    h, w = logo_alpha.shape

    # 确保位置在图像范围内
    if x + w > background.shape[1]: w = background.shape[1] - x
    if y + h > background.shape[0]: h = background.shape[0] - y

    # 调整logo大小以适应ROI
    logo_alpha_roi = logo_alpha[:h, :w]
    colored_logo_roi = colored_logo[:h, :w]

    # 获取背景ROI
    bg_roi = background[y:y + h, x:x + w]

    # 归一化alpha通道
    alpha = logo_alpha_roi.astype(float) / 255.0
    alpha = np.dstack([alpha, alpha, alpha])  # 转换为3通道

    # 融合图像
    blended = bg_roi.astype(float) * (1 - alpha) + colored_logo_roi.astype(float) * alpha
    blended = blended.astype(np.uint8)

    # 将融合后的区域放回背景
    result = background.copy()
    result[y:y + h, x:x + w] = blended

    return blended

def test_logo_colors(logo_path, background_path):
    background = cv2.imread(background_path)
    logo = cv2.imread(logo_path, cv2.IMREAD_UNCHANGED)

    if background is None or logo is None:
        print("Error loading images")
        return

    logo_alpha = logo[:, :, 3]

    # 不同颜色选项
    colors = {
        'White': (255, 255, 255),
        'Black': (0, 0, 0),
        'Red': (0, 0, 255),
        'Green': (0, 255, 0),
        'Blue': (255, 0, 0),
        'Yellow': (0, 255, 255)
    }

    results = []
    for color_name, color in colors.items():
        colored_log = np.zeros((logo.shape[0], logo.shape[1], 3), dtype=np.uint8)
        colored_log[:] = color

        #溶合
        alpha = logo_alpha.astype(float) / 255.0
        alpha = np.dstack([alpha, alpha, alpha])  # 转换为3通道

        h,w = logo_alpha.shape
        bg_roi = background[50:50+h, 50:50+w].astype(float)

        blended = bg_roi * (1 - alpha) + colored_log.astype(float) * alpha
        blended = blended.astype(np.uint8)

        preview = background.copy()
        preview[50:50+h, 50:50+w] = blended

        #添加文字
        cv2.putText(preview, color_name, (50, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
        results.append(preview)

    for i, result in enumerate(results):
        cv2.imshow(f"Logo with {list(colors.keys())[i]}", result)
        cv2.waitKey(0)
        cv2.destroyAllWindows()


def compare_blending_methods():
    """
    对比不同混合方法的效果
    """
    # 创建测试图像
    bg = np.ones((200, 300, 3), dtype=np.uint8) * 255  # 白色背景
    fg = np.zeros((200, 300, 3), dtype=np.uint8)  # 黑色前景
    fg[50:150, 100:200] = [0, 0, 255]  # 红色矩形

    # 创建alpha通道（圆形渐变）
    alpha = np.zeros((200, 300), dtype=np.uint8)
    cv2.circle(alpha, (150, 100), 60, 255, -1)
    alpha = cv2.GaussianBlur(alpha, (25, 25), 0)
    alpha_3ch = np.dstack([alpha, alpha, alpha]).astype(float) / 255.0

    # 方法1：Alpha混合（平滑过渡）
    blended_smooth = bg.astype(float) * (1 - alpha_3ch) + fg.astype(float) * alpha_3ch
    blended_smooth = blended_smooth.astype(np.uint8)

    # 方法2：二值混合（硬边缘）
    _, binary_alpha = cv2.threshold(alpha, 127, 255, cv2.THRESH_BINARY)
    binary_alpha_3ch = np.dstack([binary_alpha, binary_alpha, binary_alpha]).astype(float) / 255.0
    blended_binary = bg.astype(float) * (1 - binary_alpha_3ch) + fg.astype(float) * binary_alpha_3ch
    blended_binary = blended_binary.astype(np.uint8)

    # 显示对比
    cv2.imshow("Alpha Channel", alpha)
    cv2.imshow("Smooth Alpha Blending", blended_smooth)
    cv2.imshow("Binary Alpha Blending", blended_binary)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

# 使用示例
# 读取背景图像
background = cv2.imread('surprise.jpg')
if background is None:
    print("Error: Could not load surprise.jpg")
else:
    # 叠加白色logo在左上角
    result = overlay_logo(background, 'logo.png', position=(50, 50), color=(255, 255, 255))

    cv2.imshow("With Logo", result)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

if __name__ == '__main__':
    test_logo_colors('logo.png', 'surprise.jpg')
    compare_blending_methods()